use std::sync::Arc;

use bytestring::ByteString;
use dashmap::DashMap;
use datafusion::arrow::array::{GenericStringBuilder, RecordBatch};
use datafusion::arrow::datatypes::{Field, Schema};
use datafusion::arrow::datatypes::DataType::Utf8;
use datafusion::dataframe::DataFrame;
use datafusion::prelude::SessionContext;
use sqlparser::ast::{Ident, ObjectName, ObjectType, ShowStatementFilter};

use common::{affect_rows_batch, Entry, Error, LResult};
use log::info;

use crate::{Engine, EngineFactoryRef, EngineRef};

pub type StorageRef = Arc<Storage>;

pub struct Storage {
    tables: DashMap<String, Box<dyn Engine>>,
    engine_factories: DashMap<String, EngineFactoryRef>,
    default_engine_name: Option<String>,
    ctx: Arc<SessionContext>,
}

impl Storage {
    pub fn default() -> Self {
        let ctx = Arc::new(SessionContext::new());
        Self::create(ctx)
    }

    pub fn create(ctx: Arc<SessionContext>) -> Self {
        Self {
            tables: Default::default(),
            engine_factories: Default::default(),
            default_engine_name: Some("mem_table".into()),
            ctx,
        }
    }
    pub async fn insert(&self, line: ByteString) -> LResult<()> {
        let entry = Entry::parse(line)?;
        if let Some(mut engine) = self.tables.get_mut(&entry.measurement) {
            engine.write(entry).await
        } else {
            let measurement = entry.measurement.clone();
            let mut engine = self.new_engine(&measurement);
            engine.write(entry).await?;
            self.ctx.register_table(engine.name(), engine.table_provider()?).unwrap();
            self.tables.insert(measurement, engine);
            Ok(())
        }
    }

    pub fn post_transform_query(&self, df: DataFrame) -> LResult<DataFrame> {
        let factory = self.engine_factories.get(self.default_engine_name.as_ref().unwrap()).unwrap();
        factory.post_transform_query(df)
    }

    fn new_engine(&self, measurement: &String) -> EngineRef {
        let factory = self.engine_factories.get(self.default_engine_name.as_ref().unwrap()).unwrap();
        factory.create(measurement.clone())
    }

    pub fn register_engine_factory(&mut self, factory: EngineFactoryRef) {
        self.engine_factories.insert(factory.name().into(), factory);
    }

    // TODO, show engine information if full
    pub async fn show_tables(&self, _extended: bool, _full: bool, db_name: &Option<Ident>, filter: &Option<ShowStatementFilter>) -> LResult<RecordBatch> {
        let schema = Schema::new(vec![Field::new("name", Utf8, false)]);
        let mut c0 = GenericStringBuilder::<i32>::new();
        for e in &self.tables {
            c0.append_value(e.key());
        }
        RecordBatch::try_new(Arc::new(schema), vec![Arc::new(c0.finish())]).map_err(|e| {
            Error::StorageErr(format!("show_tables build record batch err: {}", e))
        })
    }

    pub async fn drop(&self, object_type: &ObjectType, if_exists: bool, names: &Vec<ObjectName>) -> LResult<RecordBatch> {
        if matches!(object_type, ObjectType::Table) {
            let mut rows = 0;
            for name in names {
                if name.0.len() != 1 {
                    return Err(Error::StorageErr("Syntax error, Only support DROP TableName".into()));
                }
                let table_name = &name.0[0].value;
                if let Some((_, mut v)) = self.tables.remove(table_name) {
                    v.delete().await?;
                    rows += 1;
                } else if if_exists {
                    // do nothing
                } else {
                    return Err(Error::StorageErr(format!("Table {} not exists", table_name)));
                }
            }
            Ok(affect_rows_batch(rows))
        } else {
            Err(Error::StorageErr(format!("drop {} not support!", object_type)))
        }
    }

    pub async fn create_database(&self, db_name: &ObjectName, if_not_exists: bool) -> LResult<RecordBatch> {
        todo!()
    }
}